from typing import Annotated, Sequence
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import PromptTemplate

from langgraph.checkpoint.memory import MemorySaver
from langgraph.prebuilt import ToolNode
from typing_extensions import TypedDict
from langchain_community.vectorstores import Chroma
from langchain_core.tools import create_retriever_tool

from langchain_core.messages import BaseMessage
from langgraph.graph import add_messages, StateGraph,END,START

from qw_model import qw_model
from Embed_model import BGE_Embed

# -------------------------------------定义工具
# 加载知识库
embedding = BGE_Embed()
vectorstore = Chroma(
    embedding_function=embedding,
    persist_directory='./chroma_db'
)

# 创建检索器
retriever = vectorstore.as_retriever(search_kwargs={'k': 3})

# 创建retriever工具
retriever_tool = create_retriever_tool(
    retriever,
    "yingong",
    "业务咨询和价格查询"
)

# 工具列表
tools = [retriever_tool]

# # 检索出的文档
# result = retriever.invoke('抖音云剪多少钱')
# print(result)


# ————————————————————————————————————定义节点
# 状态图
class AgentState(TypedDict):
    # add_messages 函数定义了应如何处理更新。
    # 默认是替换。add_messages 表示 “追加”。
    messages: Annotated[Sequence[BaseMessage], add_messages]

def generate(state):
    messages = state["messages"]
    question = messages[0].content
    last_message = messages[-1]

    docs = last_message.content

    # Prompt
    prompt = PromptTemplate(
        template="""你是一个聪明又有亲和力的销售客服，名叫银弓，擅长用轻松、亲切的语气解答用户关于公司业务、价格、产品咨询的问题。你的目标是提供专业、准确的回答，尤其是价格问题，严格按照相应文档回答，风格需要口语化、自然流畅、贴近用户生活，避免书面化和过于正式的表达以及尊称。\n 
             这是检索到的文档: \n\n {context} \n\n
             这是用户问题: {question} \n
             "1. 当用户问题不清楚时，主动询问用户具体需求。"
             "2. 回答价格问题，严格按照检索的文档给出，不可自行编造。"
             "3. 当没有用户所说的地区，就按它的省份价格算，当省份也没有那就只能做不限地区的"
             "4. 保持沟通语气亲切、自然，避免过于正式的措辞，多用口语化的表达和轻松的语气，比如“老板”、“宝子”、“咱们”，禁止使用'您'。"
             "5. 输出答案应简洁明了，字数在15字左右。" \n
              需满足以上回答要求。""",
        input_variables=["context", "question"],
    )

    # LLM
    llm = qw_model()
    model = llm.bind_tools(tools)

    # Chain
    rag_chain = prompt | model | StrOutputParser()

    response = rag_chain.invoke({"context": docs, "question": question})
    return {"messages": [response]}


#---------------------------------------------初始化图形
# 定义转状态图
w = StateGraph(AgentState)

# 节点

retrieve = ToolNode([retriever_tool])
w.add_node("generate",generate)

w.add_edge(START, "generate")

# --------------------------------------------编译图
memory = MemorySaver()
app = w.compile(checkpointer=memory)

config = {"configurable": {'thread_id': "a1"}}

# # 可视化
# graph_img = app.get_graph().draw_mermaid_png(
#     draw_method=MermaidDrawMethod.API,
# )
# with open("graph_1.png", "wb") as f:
#     f.write(graph_img)


res = app.invoke({
    "messages": [
        ("user", "实探多少钱"),
    ]
}, config)
print(res["messages"][-1].content)

res2 = app.invoke({
    "messages": [
        ("user", "成都，最好是千粉的"),
    ]
}, config)
print(res2["messages"][-1].content)
